import torch
import torch.nn as nn
import numpy as np

from model.TemperatureHumidityMLP import TemperatureHumidityMLP

# 初始化模型
input_size = 7
hidden_size1 = 64
hidden_size2 = 64
num_classes_temp = 42 - 30 + 1
num_classes_humidity = 60 - 40 + 1
model = TemperatureHumidityMLP(input_size, hidden_size1, hidden_size2, num_classes_temp, num_classes_humidity)

# 加载模型参数
model.load_state_dict(torch.load('ckpt/model.pth', map_location='cpu'))
model.eval()

# 导出模型参数到 .txt 文件
def export_weights_to_c_file(model, filename):
    with open(filename, 'w') as f:
        # 导出 fc1 层的权重和偏置
        fc1_weights = model.fc1.weight.detach().cpu().numpy()
        fc1_bias = model.fc1.bias.detach().cpu().numpy()
        f.write(f"float fc1_weights[{hidden_size1}][{input_size}] = {{\n")
        for row in fc1_weights:
            f.write("    {" + ", ".join([f"{w:.6f}f" for w in row]) + "},\n")
        f.write("};\n\n")
        f.write(f"float fc1_bias[{hidden_size1}] = {{\n")
        f.write(", ".join([f"{b:.6f}f" for b in fc1_bias]))
        f.write("};\n\n")

        # 导出 fc2 层的权重和偏置
        fc2_weights = model.fc2.weight.detach().cpu().numpy()
        fc2_bias = model.fc2.bias.detach().cpu().numpy()
        f.write(f"float fc2_weights[{hidden_size2}][{hidden_size1}] = {{\n")
        for row in fc2_weights:
            f.write("    {" + ", ".join([f"{w:.6f}f" for w in row]) + "},\n")
        f.write("};\n\n")
        f.write(f"float fc2_bias[{hidden_size2}] = {{\n")
        f.write(", ".join([f"{b:.6f}f" for b in fc2_bias]))
        f.write("};\n\n")

        # 导出 fc3_temp 层的权重和偏置
        fc3_temp_weights = model.fc3_temp.weight.detach().cpu().numpy()
        fc3_temp_bias = model.fc3_temp.bias.detach().cpu().numpy()
        f.write(f"float fc3_temp_weights[{num_classes_temp}][{hidden_size2}] = {{\n")
        for row in fc3_temp_weights:
            f.write("    {" + ", ".join([f"{w:.6f}f" for w in row]) + "},\n")
        f.write("};\n\n")
        f.write(f"float fc3_temp_bias[{num_classes_temp}] = {{\n")
        f.write(", ".join([f"{b:.6f}f" for b in fc3_temp_bias]))
        f.write("};\n\n")

        # 导出 fc3_humidity 层的权重和偏置
        fc3_humidity_weights = model.fc3_humidity.weight.detach().cpu().numpy()
        fc3_humidity_bias = model.fc3_humidity.bias.detach().cpu().numpy()
        f.write(f"float fc3_humidity_weights[{num_classes_humidity}][{hidden_size2}] = {{\n")
        for row in fc3_humidity_weights:
            f.write("    {" + ", ".join([f"{w:.6f}f" for w in row]) + "},\n")
        f.write("};\n\n")
        f.write(f"float fc3_humidity_bias[{num_classes_humidity}] = {{\n")
        f.write(", ".join([f"{b:.6f}f" for b in fc3_humidity_bias]))
        f.write("};\n")

# 调用导出函数
export_weights_to_c_file(model, 'model_parameters.c')
print("Model parameters exported to model_parameters.c")
